"""
Unary-pairwise transformer for human-object interaction detection

Fred Zhang <frederic.zhang@anu.edu.au>

The Australian National University
Australian Centre for Robotic Vision
"""

import os
import torch
import torch.distributed as dist


from torch import nn, Tensor
from typing import Optional, List
from torchvision.ops.boxes import batched_nms, box_iou

from models.upt_fasterrcnn.interaction_head import InteractionHead
from models.upt_fasterrcnn.faster_rcnn import fasterrcnn_resnet_fpn
from utils.misc import nested_tensor_from_tensor_list

from typing import Optional, List, Tuple
from torchvision.models.detection import transform

class UPT_Fasterrcnn(nn.Module):
    """
    Unary-pairwise transformer

    Parameters:
    -----------
    detector: nn.Module
        Object detector (DETR)
    postprocessor: nn.Module
        Postprocessor for the object detector
    interaction_head: nn.Module
        Interaction head of the network
    human_idx: int
        Index of the human class
    """
    def __init__(self, cfg,
        backbone: nn.Module,
        interaction_head: nn.Module,
    ) -> None:
        super().__init__()
        self.cfg = cfg
        self.backbone = backbone
        self.interaction_head = interaction_head

        self.human_idx = cfg.DATASET.HUMAN_IDX
        self.num_classes = cfg.DATASET.NUM_CLASSES

        self.box_score_thresh = cfg.DATASET.BOX_SCORE_THRESH
        self.fg_iou_thresh = cfg.DATASET.FG_IOU_THRESH

        self.min_instances = cfg.DATASET.MIN_INSTANCES
        self.max_instances = cfg.DATASET.MAX_INSTANCES

    def prepare_region_proposals(self, results, hidden_states):
        region_props = []
        for res, hs in zip(results, hidden_states):
            sc, lb, bx = res.values()

            keep = batched_nms(bx, sc, lb, 0.5)
            sc = sc[keep].view(-1)
            lb = lb[keep].view(-1)
            bx = bx[keep].view(-1, 4)
            hs = hs[keep].view(-1, 256)

            keep = torch.nonzero(sc >= self.box_score_thresh).squeeze(1)

            is_human = lb == self.human_idx
            hum = torch.nonzero(is_human).squeeze(1)
            obj = torch.nonzero(is_human == 0).squeeze(1)
            n_human = is_human[keep].sum(); n_object = len(keep) - n_human
            # Keep the number of human and object instances in a specified interval
            if n_human < self.min_instances:
                keep_h = sc[hum].argsort(descending=True)[:self.min_instances]
                keep_h = hum[keep_h]
            elif n_human > self.max_instances:
                keep_h = sc[hum].argsort(descending=True)[:self.max_instances]
                keep_h = hum[keep_h]
            else:
                keep_h = torch.nonzero(is_human[keep]).squeeze(1)
                keep_h = keep[keep_h]

            if n_object < self.min_instances:
                keep_o = sc[obj].argsort(descending=True)[:self.min_instances]
                keep_o = obj[keep_o]
            elif n_object > self.max_instances:
                keep_o = sc[obj].argsort(descending=True)[:self.max_instances]
                keep_o = obj[keep_o]
            else:
                keep_o = torch.nonzero(is_human[keep] == 0).squeeze(1)
                keep_o = keep[keep_o]

            keep = torch.cat([keep_h, keep_o])

            region_props.append(dict(
                boxes=bx[keep],
                scores=sc[keep],
                labels=lb[keep],
                hidden_states=hs[keep]
            ))

        return region_props

    def postprocessing(self, boxes, bh, bo, logits, prior, objects, attn_maps, image_sizes):
        n = [len(b) for b in bh]
        logits = logits.split(n)

        detections = []
        for bx, h, o, lg, pr, obj, attn, size in zip(
            boxes, bh, bo, logits, prior, objects, attn_maps, image_sizes
        ):
            pr = pr.prod(0)
            x, y = torch.nonzero(pr).unbind(1)
            scores = torch.sigmoid(lg[x, y])
            detections.append(dict(
                boxes=bx, pairing=torch.stack([h[x], o[x]]),
                scores=scores * pr[x, y], labels=y,
                objects=obj[x], attn_maps=attn, size=size
            ))

        return detections

    def preprocess_boxes(self,
        targets: List[dict],
        append_gt: Optional[bool] = None
    ) -> None:
        results = []
        if self.cfg.MODEL.ONLY_GT_BOX:
            for target in targets:
                n = target["boxes_h"].shape[0]
                device = target["boxes_h"].device
                boxes = torch.cat([target["boxes_h"], target["boxes_o"]])
                scores = torch.cat([torch.ones(2 * n, device=device)])
                labels = torch.cat([
                        self.human_idx * torch.ones(n, device=device).long(),
                        target["object"]
                    ])

                # Keep a fixed number of detections
                h_idx = torch.nonzero(labels == self.human_idx).squeeze(1)
                o_idx = torch.nonzero(labels != self.human_idx).squeeze(1)
                if len(h_idx) > self.max_instances:
                    h_idx = h_idx[:self.max_instances]
                if len(o_idx) > self.max_instances:
                    o_idx = o_idx[:self.max_instances]
                # Permute humans to the top
                keep_idx = torch.cat([h_idx, o_idx])
                results.append(dict(
                    boxes=boxes[keep_idx].view(-1, 4),
                    labels=labels[keep_idx].view(-1),
                    scores=scores[keep_idx].view(-1)
                ))
        else:
            for b_idx, target in enumerate(targets):
                detection = target['detections']
                boxes = detection['boxes']
                labels = detection['labels']
                scores = detection['scores']

                # Append ground truth during training
                if append_gt is None:
                    append_gt = self.training
                if append_gt:
                    n = target["boxes_h"].shape[0]
                    boxes = torch.cat([target["boxes_h"], target["boxes_o"], boxes])
                    scores = torch.cat([torch.ones(2 * n, device=scores.device), scores])
                    labels = torch.cat([
                        self.human_idx * torch.ones(n, device=labels.device).long(),
                        target["object"],
                        labels
                    ])

                # Remove low scoring examples
                active_idx = torch.nonzero(
                    scores >= self.box_score_thresh
                ).squeeze(1)
                box_nms_thresh = 0.5
                
                # Class-wise non-maximum suppression
                keep_idx = batched_nms(
                    boxes[active_idx],
                    scores[active_idx],
                    labels[active_idx],
                    box_nms_thresh
                )
                active_idx = active_idx[keep_idx]
                # Sort detections by scores
                sorted_idx = torch.argsort(scores[active_idx], descending=True)
                active_idx = active_idx[sorted_idx]
                # Keep a fixed number of detections
                h_idx = torch.nonzero(labels[active_idx] == self.human_idx).squeeze(1)
                o_idx = torch.nonzero(labels[active_idx] != self.human_idx).squeeze(1)
                if len(h_idx) > self.max_instances:
                    h_idx = h_idx[:self.max_instances]
                if len(o_idx) > self.max_instances:
                    o_idx = o_idx[:self.max_instances]
                # Permute humans to the top
                keep_idx = torch.cat([h_idx, o_idx])
                active_idx = active_idx[keep_idx]

                results.append(dict(
                    boxes=boxes[active_idx].view(-1, 4),
                    labels=labels[active_idx].view(-1),
                    scores=scores[active_idx].view(-1)
                ))

        return results
    
    def forward(self,
        images: List[Tensor],
        targets: Optional[List[dict]] = None,
        is_train: bool = True
    ) -> List[dict]:
        """
        Parameters:
        -----------
        images: List[Tensor]
            Input images in format (C, H, W)
        targets: List[dict], optional
            Human-object interaction targets

        Returns:
        --------
        results: List[dict]
            Detected human-object interactions. Each dict has the following keys:
            `boxes`: torch.Tensor
                (N, 4) Bounding boxes for detected human and object instances
            `pairing`: torch.Tensor
                (2, M) Pairing indices, with human instance preceding the object instance
            `scores`: torch.Tensor
                (M,) Interaction score for each pair
            `labels`: torch.Tensor
                (M,) Predicted action class for each pair
            `objects`: torch.Tensor
                (M,) Predicted object class for each pair
            `attn_maps`: list
                Attention weights in the cooperative and competitive layers
            `size`: torch.Tensor
                (2,) Image height and width
        """
        if self.training and targets is None:
            raise ValueError("In training mode, targets should be passed")
        image_sizes = torch.as_tensor([
            im.size()[-2:] for im in images
        ], device=images[0].device)

        if isinstance(images, (list, torch.Tensor)):
            images = nested_tensor_from_tensor_list(images)
        features = self.backbone(images.tensors)
        
        region_props = self.preprocess_boxes(targets)

        logits, prior, bh, bo, objects, attn_maps = self.interaction_head(
            features, image_sizes, region_props
        )
        boxes = [r['boxes'] for r in region_props]

        return boxes, bh, bo, logits, prior, objects, attn_maps, image_sizes

        if self.training:
            interaction_loss = self.compute_interaction_loss(boxes, bh, bo, logits, prior, targets)
            loss_dict = dict(
                interaction_loss=interaction_loss
            )
            return loss_dict

        detections = self.postprocessing(boxes, bh, bo, logits, prior, objects, attn_maps, image_sizes)
        return detections

def build_upt_facter_rcnn(cfg, args, class_corr):
    
    detector = fasterrcnn_resnet_fpn(cfg.MODEL.BACKBONE, pretrained=True)
    backbone = detector.backbone
    
    predictor = torch.nn.Linear(cfg.MODEL.REPR_DIM * 2, cfg.DATASET.NUM_CLASSES)
    interaction_head = InteractionHead(
        predictor, cfg.MODEL.HIDDEN_DIM, cfg.MODEL.REPR_DIM,
        256,    # faster-rcnn输出的dim=256
        cfg.DATASET.NUM_CLASSES, cfg.DATASET.HUMAN_IDX, class_corr
    )
    model = UPT_Fasterrcnn(cfg,
        backbone, interaction_head,
    )
    return model
